import yaml
import yaml.parser
from core.status_logger import LogStatus as CLogger


class YamlLoaderWithLog():
    def __init__(self, file=None, stream=None, encoding="utf-8", loader=yaml.SafeLoader, logger: CLogger=None, raise_again=False):
        self._yaml_dict = {}
        self._file = file
        self._stream = stream
        self._encoding = encoding
        self._loader = loader
        self._logger = logger
        self._raise_again = raise_again
        pass

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        if exc_val is None:
            return not self._raise_again
        try:
            if self._logger:
                self._logger.line_feed()
                if issubclass(exc_type, yaml.MarkedYAMLError):
                    self._logger.write(exc_val.problem, CLogger.MsgType.ERROR)
                    self._logger.line_feed()
                    if self._file:
                        self._logger.write("{} #(line:{}, col:{})".format(exc_val.problem_mark.name,
                                                                          exc_val.problem_mark.line,
                                                                          exc_val.problem_mark.column),
                                           CLogger.MsgType.ERROR)
                        self._logger.line_feed()
                        try:
                            with open(exc_val.problem_mark.name, 'r', encoding=self._encoding) as f:
                                lines = [next(f) for _ in range(exc_val.problem_mark.line+1)]
                                self._logger.write("  >>> {}".format(lines[-1]), CLogger.MsgType.ERROR)
                        except:
                            pass
                    elif self._stream:
                        lines = exc_val.problem_mark.buffer.splitlines()
                        self._logger.write("  >>> {}".format(lines[exc_val.problem_mark.line]),
                                           CLogger.MsgType.ERROR)
                    else:
                        pass
                elif issubclass(exc_type, OSError):
                    self._logger.write("{} ({})".format(exc_val.strerror,
                                                        exc_val.filename),
                                       CLogger.MsgType.ERROR)
                else:
                    self._logger.write(str(exc_val), CLogger.MsgType.ERROR)
                self._logger.line_feed()
            return not self._raise_again
        except Exception as new_exc:
            return False

    def load(self, file=None, stream=None, multi=False):
        if file:
            self._file = file
        if stream:
            self._stream = stream
        yaml_load = yaml.load_all if multi else yaml.load
        if self._file:
            with open(self._file, "r", encoding=self._encoding) as f:
                self._yaml_dict = yaml_load(f, Loader=self._loader)
                return self._yaml_dict
        elif self._stream:
            self._yaml_dict = yaml_load(self._stream, Loader=self._loader)
            return self._yaml_dict
        else:
            return [] if multi else {}

    def safe_load(self, file=None, stream=None, multi=False):
        try:
            return self.load(file=file, stream=stream, multi=multi)
        except Exception as ex:
            if not self.__exit__(type(ex), ex, None):
                raise ex
            return [] if multi else {}

    def load_dict(self):
        d = self.load(False)
        if isinstance(d, dict): # single doc
            return d
        else: # bad
            return {}

    def load_list(self):
        d = self.load(True)
        if isinstance(d, dict): # single doc
            return [d]
        elif isinstance(d, (list, tuple)): # multi doc
            return list(d)
        else: # bad
            return []

class YamlDumperWithLog():
    def __init__(self, file=None, encoding="utf-8", dumper=yaml.SafeDumper, logger: CLogger=None, raise_again=False):
        self._file = file
        self._encoding = encoding
        self._dumper = dumper
        self._logger = logger
        self._raise_again = raise_again
        pass

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        if exc_val is None:
            return not self._raise_again
        try:
            if self._logger:
                self._logger.line_feed()
                if issubclass(exc_type, OSError):
                    self._logger.write("{} ({})".format(exc_val.strerror,
                                                        exc_val.filename),
                                       CLogger.MsgType.ERROR)
                else:
                    self._logger.write(str(exc_val), CLogger.MsgType.ERROR)
                self._logger.line_feed()
            return not self._raise_again
        except Exception as new_exc:
            return False

    def dump(self, data, file=None, multi=False, **kwargs):
        if file:
            self._file = file
        yaml_dump = yaml.dump_all if multi else yaml.dump
        if self._file:
            with open(self._file, "w+", encoding=self._encoding) as f:
                return yaml_dump(data, stream=f, Dumper=self._dumper, **kwargs)
        else:
            return yaml_dump(data, Dumper=self._dumper, **kwargs)

    def safe_dump(self, data, file=None, multi=False, **kwargs):
        try:
            return self.dump(data, file=file, multi=multi, **kwargs)
        except Exception as ex:
            if not self.__exit__(type(ex), ex, None):
                raise ex
            return False
